#include "componentMatch.h"

void componentMatch(pcl::PointCloud<pcl::PointXYZRGB>::Ptr registeredSource, Component componentToFind, pcl::PointIndices::Ptr sourceComponentIndices, float radius)
{

    vector<bool> processed (registeredSource->points.size(), false);
    vector<int> nn_indices;
    vector<float> nn_distances;
    vector<int> seed_queue;
    int sq_idx = 0;
    vector<int> indices;
    vector<float> sqr_distances;
    KdTreeFLANN<PointXYZRGB> tree;
    tree.setInputCloud(registeredSource);
    pcl::PointIndices::Ptr unfilteredSourceComponentIndices(new PointIndices);;
    //prepare starting data
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr componentCloud = componentToFind.getComponentCloud();
    for (int i = 0; i < componentCloud->size(); i++) {
        tree.radiusSearch(componentCloud->at(i), radius*0.5, indices, sqr_distances);
        for (int j=0; j < indices.size(); j++)
            seed_queue.push_back(indices.at(j));
    }
    sort(seed_queue.begin (), seed_queue.end ());
    seed_queue.erase(unique(seed_queue.begin(), seed_queue.end()), seed_queue.end());
    int r(0),g(0),b(0);
    pcl::PointIndices::Ptr coreIndices(new PointIndices);
    coreIndices->indices.resize (seed_queue.size());
    for (size_t j = 0; j < seed_queue.size (); ++j)
        coreIndices->indices[j] = seed_queue[j];
    getAverageColor(registeredSource, coreIndices, r, g ,b);

    //end preparation

    while (sq_idx < static_cast<int> (seed_queue.size()) ){
        // Search for sq_idx
        if (!tree.radiusSearch(seed_queue[sq_idx], radius, nn_indices, nn_distances)){
            sq_idx++;
            continue;
        }
        for (size_t j = 0; j < nn_indices.size (); ++j){ // nn_indices[0] should be sq_idx
            if (nn_indices[j] == -1 || processed[nn_indices[j]]) // Has this point been processed before ?
                continue;

            // Perform a simple Euclidean clustering
            seed_queue.push_back(nn_indices[j]);
            processed[nn_indices[j]] = true;
        }

        sq_idx++;
    }

    unfilteredSourceComponentIndices->indices.resize (seed_queue.size());
    for (size_t j = 0; j < seed_queue.size (); ++j)
        unfilteredSourceComponentIndices->indices[j] = seed_queue[j];
    removeColorOutliers(registeredSource, unfilteredSourceComponentIndices, sourceComponentIndices, r, g, b, componentToFind.getColorThreshold());
}

void getAverageColor
(
        pcl::PointCloud<pcl::PointXYZRGB>::Ptr registeredSource,
        pcl::PointIndices::Ptr sourceComponentIndices,
        int &averageR,
        int &averageG,
        int &averageB
        )
{
    pcl::PointCloud<pcl::PointXYZRGB>::Ptr sourceClusterCore(new PointCloud<PointXYZRGB>);
    copyPointCloud(*registeredSource, *sourceComponentIndices, *sourceClusterCore);
    int r(0), g(0), b(0);
    if (sourceClusterCore->points.size() == 0)
    {
        averageR = 0;
        averageG = 0;
        averageB = 0;
    }
    else
    {
        for (int i = 0; i < sourceClusterCore->points.size(); i++)
        {
            r = r + sourceClusterCore->points.at(i).r;
            g = g + sourceClusterCore->points.at(i).g;
            b = b + sourceClusterCore->points.at(i).b;
        }
        averageR = r/sourceClusterCore->points.size();
        averageG = g/sourceClusterCore->points.size();
        averageB = b/sourceClusterCore->points.size();
    }
}

void removeColorOutliers
(
        pcl::PointCloud<pcl::PointXYZRGB>::Ptr cloud,
        pcl::PointIndices::Ptr input,
        pcl::PointIndices::Ptr output,
        int r, int g, int b,
        int threshold
        )
{
    for (int i = 0; i < input->indices.size(); i++)
    {
        int currentPointIndex = input->indices.at(i);
        if (
                abs(cloud->points.at(currentPointIndex).r-r) < threshold &&
                abs(cloud->points.at(currentPointIndex).g-g) < threshold &&
                abs(cloud->points.at(currentPointIndex).b-b) < threshold
                )
            output->indices.push_back(currentPointIndex);
    }

}

